import torch
import torch.nn as nn
import torch.optim as optim

# Define model
model = torch.nn.Sequential(
  ResNet(), 
  nn.Linear(512, 2)
)

# Loss and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001)

# Prepare data
# X contains 10000 128-d samples of class 0 and 10000 of class 1
X = torch.from_numpy(data)  
y = torch.from_numpy(labels) 

# Train the model
num_epochs = 10
for epoch in range(num_epochs):

  # Forward pass
  outputs = model(X)
  
  # Compute loss 
  loss = criterion(outputs, y)
  
  # Backward and optimize
  optimizer.zero_grad()
  loss.backward()
  optimizer.step()
  
print(f'Loss after epoch {epoch}: {loss.item():.4f}')